from vtkmodules.vtkRenderingCore import vtkProperty, vtkShaderProperty
from vtkmodules.vtkImagingCore import vtkImageExtractComponents, vtkImageAppendComponents

from vtkmodules.vtkRenderingCore import vtkTexture

class Material(vtkProperty):
    # load from gltf material...
    def __init__(self, mtl, textures) -> None:
        super().__init__()

        self.SetMaterialName(mtl.name)
        # pbr
        pbr = mtl.pbrMetallicRoughness

        if pbr != None:            
            self.SetInterpolationToPBR()
            self.SetColor(pbr.baseColorFactor[:3])
            self.SetOpacity(pbr.baseColorFactor[3])
            self.SetMetallic(pbr.metallicFactor)
            self.SetRoughness(pbr.roughnessFactor)

            if pbr.baseColorTexture:
                albedo = textures[pbr.baseColorTexture.index]
                albedo.UseSRGBColorSpaceOn()
                self.SetBaseColorTexture(albedo)

            if mtl.normalTexture != None:
                normal = textures[mtl.normalTexture.index]
                nScale = mtl.normalTexture.scale
                self.SetNormalTexture(normal)
                self.SetNormalScale(nScale)

            imgO  = None
            imgRM = None
            if pbr.metallicRoughnessTexture != None:
                mrTex = textures[pbr.metallicRoughnessTexture.index]
                filter = vtkImageExtractComponents()
                filter.SetComponents(1, 0)  # roughness, metallic
                filter.SetInputData(mrTex.GetInput())
                filter.Update()
                imgRM = filter.GetOutput()

            if mtl.occlusionTexture != None:
                oTex = textures[mtl.occlusionTexture.index]
                filter = vtkImageExtractComponents()
                filter.SetComponents(0)
                filter.SetInputData(oTex.GetInput())
                filter.Update()
                imgO = filter.GetOutput()

            if imgO and imgRM:
                filter = vtkImageAppendComponents()
                filter.SetInputData(0, imgO)
                filter.AddInputData(0, imgRM)
                filter.Update()
                imgORM = filter.GetOutput()
                texORM = vtkTexture()
                texORM.SetInputData(imgORM)
                texORM.SetInterpolate(1)
                self.SetORMTexture(texORM)

        else:
            print("Unsupported material:", mtl.name)
        
        self.shaderProperty = vtkShaderProperty()
        self._setupShaderProperty()


    # def ResetColorMode(self):
    #     self._resetShaderProperty()

    # def SetColorModeToNormal(self): 
    #     self.ResetColorMode() 
    #     self.shaderProperty.AddFragmentShaderReplacement(
    #     "//VTK::Light::Impl",  True,
    #     "//VTK::Light::Impl\n" 
    #     "gl_FragData[0] = vec4(N * 0.5 + 0.5, 1.0);\n", False)

    # def SetColorModeToRandomPrimitiveColor(self):    
    #     self.ResetColorMode() 
    #     self.shaderProperty.AddFragmentShaderReplacement(
    #     "//VTK::Color::Dec",  True,
    #     "//VTK::Color::Dec\n" 
    #     "uniform vec3 randomColorUniform;\n", False)
    #     self.shaderProperty.AddFragmentShaderReplacement(
    #     "//VTK::Light::Impl",  True,
    #     "//VTK::Light::Impl\n" 
    #     "gl_FragData[0] = vec4(randomColorUniform, 1);\n", 
    #     False)   
    
    # def SetColorModeToDepth(self, scale:float=0.1):    
    #     self.ResetColorMode() 
    #     self.shaderProperty.AddFragmentShaderReplacement(
    #     "//VTK::Light::Impl",  True,
    #     "//VTK::Light::Impl\n" 
    #     "gl_FragData[0] = vec4(vec3(abs(vertexVC.z * {})), 1);\n".format(scale), False) 
    
    # def SetColorModeToAlbedo(self):    
    #     self.ResetColorMode() 
    #     self.shaderProperty.AddFragmentShaderReplacement(
    #     "//VTK::Light::Impl",  True,
    #     "//VTK::Light::Impl\n" 
    #     "gl_FragData[0] = vec4(albedo, 1);\n", False) 


    def _setupShaderProperty(self):
        self.shaderProperty.ClearAllShaderReplacements()

        if self.GetTexture("normalTex") is not None:
            # patch for fragment tangent if vtkGlyph3DMapper used...
            self.shaderProperty.AddFragmentShaderReplacement(
            "//VTK::Edges::Impl",  True,
            "//VTK::Edges::Impl\n"         
            "vec3 Q1  = dFdx(vertexVCVSOutput.xyz);"
            "vec3 Q2  = dFdy(vertexVCVSOutput.xyz);"
            "vec2 st1 = dFdx(tcoordVCVSOutput);"
            "vec2 st2 = dFdy(tcoordVCVSOutput);"
            "vec3 tangentVCVSOutput = normalize(Q1 * st2.t - Q2 * st1.t);\n", 
            False)

        # for debug
        # self.shaderProperty.AddFragmentShaderReplacement(
        #     "//VTK::Light::Impl",  True,
        #     "//VTK::Light::Impl\n"         
        #     "fragOutput0 = vec4(vec3(ao), opacity);;\n", 
        #     False)
            
    def GetShaderProperty(self):
        sp = vtkShaderProperty()
        sp.DeepCopy(self.shaderProperty)
        return sp